package com.kdpujie.alink;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.classification.LogisticRegressionPredictBatchOp;
import com.alibaba.alink.operator.batch.classification.LogisticRegressionTrainBatchOp;
import com.alibaba.alink.operator.batch.dataproc.SplitBatchOp;
import com.alibaba.alink.operator.batch.evaluation.EvalBinaryClassBatchOp;
import com.alibaba.alink.operator.batch.source.CsvSourceBatchOp;
import com.alibaba.alink.operator.common.evaluation.BinaryClassMetrics;
import com.alibaba.alink.pipeline.Pipeline;
import com.alibaba.alink.pipeline.dataproc.StandardScaler;
import com.alibaba.alink.pipeline.feature.FeatureHasher;
import com.kdpujie.alink.source.AvazuCtrSource;
import com.kdpujie.alink.source.SSPCtrSource;
import com.kdpujie.alink.source.Source;

/**
 * 特征工程、训练、评估LR模型
 */
public class LrClassifier {
    public static void main(String[] args) {
        Source<CsvSourceBatchOp> source = new SSPCtrSource<>();

        BatchOperator<?> originTrainData = new SplitBatchOp().setFraction(0.9).linkFrom(source.GetSource());
        BatchOperator<?> originPredictData = originTrainData.getSideOutput(0); 

        // 特征工程后，把多特征合并成一个特征向量，该特征向量的列名
        String vectorlName = "vec";
        // 特征工程的 pipeline
        Pipeline featurePipeline = new Pipeline()
        .add(
            new FeatureHasher().setSelectedCols(source.GetColsName()).setCategoricalCols(source.GetColsName()).setOutputCol(vectorlName).setNumFeatures(30000)
        ); 
        //特征工程pipeline持久化
        //featurePipeline.fit(batchData).save("/Users/pujie/codes/java-projects/practise-flink/alink-ftrl/target/feature_pipe_model.csv");
        try {
            // originBatchData.print();
            // 数值标准化
            // BatchOperator<?> standardScaler = new StandardScalerTrainBatchOp().setSelectedCols(numericalNames);
            // standardScaler.linkFrom(batchData);
            // BatchOperator<?> predictOp = new StandardScalerPredictBatchOp();
            // predictOp.linkFrom(standardScaler, batchData);
            // // FeatureHasher
            // BatchOperator<?> haser = new FeatureHasherBatchOp().setSelectedCols(selectedNames).setOutputCol("vec").setNumFeatures(30000);
            // haser.linkFrom(predictOp).print();
            //特征工程
            BatchOperator<?> trainData = featurePipeline.fit(originTrainData).transform(originTrainData);
            BatchOperator<?> predictData = featurePipeline.fit(originPredictData).transform(originPredictData);

            // LR模型
            LogisticRegressionTrainBatchOp lr = new LogisticRegressionTrainBatchOp().setVectorCol(vectorlName).setLabelCol(source.GetLabelName()).setWithIntercept(true).setMaxIter(100);
            BatchOperator<?> initModel = trainData.link(lr); //
            
            // 预测
            BatchOperator<?> predictor1 = new LogisticRegressionPredictBatchOp().setPredictionCol("pred").setPredictionDetailCol("detail").linkFrom(initModel, trainData);
            BatchOperator<?> predictor2 = new LogisticRegressionPredictBatchOp().setPredictionCol("pred").setPredictionDetailCol("detail").linkFrom(initModel, predictData);

            // predictor.linkFrom(initModel, trainData).sample(0.001).print();

            // 评估
            // 对训练数据的评估
            BinaryClassMetrics metrics1 = new EvalBinaryClassBatchOp().setLabelCol(source.GetLabelName()).setPredictionDetailCol("detail").linkFrom(predictor1).collectMetrics();
            System.out.println("对训练数据的评估：");
            System.out.println("\tAUC:" + metrics1.getAuc());
            System.out.println("\tAccuracy:" + metrics1.getAccuracy());
            System.out.println("\tLogLoss:" + metrics1.getLogLoss());

            // 对测试数据的评估
            BinaryClassMetrics metrics2 = new EvalBinaryClassBatchOp().setLabelCol(source.GetLabelName()).setPredictionDetailCol("detail").linkFrom(predictor2).collectMetrics();
            System.out.println("对测试数据的评估：");
            System.out.println("\tAUC:" + metrics2.getAuc());
            System.out.println("\tAccuracy:" + metrics2.getAccuracy());
            System.out.println("\tLogLoss:" + metrics2.getLogLoss());

        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}